#! /usr/bin/env python
#coding=utf-8

import os
import codecs

from smaps import VirtualMemroyArea
from smaps import ProcessSmaps
from threads import parse_threads_file

class ProcessPsParser(dict):
	def __init__(self, parts):
		self["user"] = parts[0]
		self["pid"] = int(parts[1])
		self["ppid"] = int(parts[2])
		self["name"] = parts[7]

class SystemObject(dict):
	def __init__(self, idx, obj):
		self["id"] = idx
		self["processes"] = 1
		for k in obj:
			self[k] = obj[k]

	def add_obj(self, obj):
		for k in VirtualMemroyArea.VMA_SIZE_KEYS:
			self[k] = self[k] + obj[k]

		self["processes"] = self["processes"] + 1

class ProcessMgr(object):
	def __init__(self, dir):
		self._procs = []
		self._objects = []
		self._objects_dict = {}

		# Check if it's new aln structure or old structure
		if os.path.exists(os.path.join(dir, "ps/ps-ef.txt")):
			# Old structure
			ps_file = os.path.join(dir, "ps/ps-ef.txt")
			with open(ps_file) as f:
				self.__load_old_structure(f, dir)
		else:
			# New aln structure
			self.__load_aln_structure(dir)

	def get_all(self):
		return self._procs

	def get_all_objects(self):
		return self._objects

	def get_object_by_name(self, name):
		if name in self._objects_dict:
			return self._objects_dict[name]
		return None

	def _merge_objects(self):
		cnt = 0
		for proc in self._procs:
			for obj in proc.get_all():
				sobj = self.get_object_by_name(obj["name"])
				if sobj:
					sobj.add_obj(obj)
					continue
				cnt = cnt + 1
				sobj = SystemObject(cnt, obj)
				self._objects.append(sobj)
				self._objects_dict[obj["name"]] = sobj

	def __load_aln_structure(self, dir):
		# Load each process directory in aln structure
		for item in os.listdir(dir):
			item_path = os.path.join(dir, item)
			if not os.path.isdir(item_path):
				continue

			# Parse directory name to extract PID and process name
			# Format: {pid}_{name}
			if "_" not in item:
				continue

			parts = item.split("_", 1)
			if len(parts) != 2:
				continue

			try:
				pid = int(parts[0])
				name = parts[1]
			except ValueError:
				continue

			# Find the latest ps and smaps files
			ps_dir = os.path.join(item_path, "ps")
			smaps_dir = os.path.join(item_path, "smaps")

			if not os.path.exists(ps_dir) or not os.path.exists(smaps_dir):
				continue

			# Get the latest timestamp file
			ps_files = [f for f in os.listdir(ps_dir) if f.startswith("ps_")]
			smaps_files = [f for f in os.listdir(smaps_dir) if f.startswith("smaps_")]

			if not ps_files or not smaps_files:
				continue

			# Sort by timestamp to get the latest
			ps_files.sort()
			smaps_files.sort()

			latest_ps_file = os.path.join(ps_dir, ps_files[-1])
			latest_smaps_file = os.path.join(smaps_dir, smaps_files[-1])

			print("Loading process %d (%s) from aln structure" % (pid, name))

			# Create process object
			proc_info = {
				"user": "unknown",
				"pid": pid,
				"ppid": 0,
				"name": name
			}

			# Load smaps info
			proc_with_smaps = ProcessSmaps(latest_smaps_file)
			for k, v in proc_info.items():
				proc_with_smaps[k] = v

			# Parse threads from ps file (since no separate threads file)
			proc_with_smaps["threads"] = self._parse_threads_from_ps(latest_ps_file)

			self._procs.append(proc_with_smaps)

		self._merge_objects()

	def _parse_threads_from_ps(self, ps_file):
		"""Parse thread information from ps file (aln structure)"""
		threads = []
		try:
			with open(ps_file, 'r') as f:
				lines = f.readlines()
				for line in lines[1:]:  # Skip header
					line = line.strip()
					if not line:
						continue
					parts = line.split()
					if len(parts) >= 5:
						# Format: PID TID TTY TIME CMD
						try:
							tid = int(parts[1])
							name = " ".join(parts[4:])  # Handle commands with spaces
							threads.append({
								"tid": tid,
								"name": name
							})
						except ValueError:
							continue
		except Exception as e:
			print("Error parsing threads from ps file %s: %s" % (ps_file, str(e)))

		return threads

	def __load_old_structure(self, f, dir):
		skip_ppids = [2] # Skip kernel process
		for line in f.readlines():
			line = line.strip()
			if "shell cat" in line:
				continue
			parts = line.split()
			if len(parts) < 8:
				continue
			if parts[0] == "UID":
				continue
			proc = ProcessPsParser(parts)
			if proc["pid"] == 2: # Skip kernel root process
				continue
			if proc["ppid"] in skip_ppids:
				#print("Skip " + proc)
				continue
			if proc["name"] == "hdcd":
				skip_ppids.append(proc["pid"])

			# Load smaps info
			smaps_file = os.path.join(dir, "pids/smaps/", "%d_smaps.txt" % proc["pid"])
			print("Loading smaps info for " + smaps_file)
			proc_with_smaps = ProcessSmaps(smaps_file)
			for k, v in proc.items():
				proc_with_smaps[k] = v

			threads_file = os.path.join(dir, "pids/threads/", "%d_threads.txt" % proc["pid"])
			proc_with_smaps["threads"] = parse_threads_file(threads_file)

			self._procs.append(proc_with_smaps)
			#print(self._procs)

		self._merge_objects()

if __name__ == "__main__":
	mgr = ProcessMgr("./aln")

	print(mgr.get_all_objects())
